local ae = require "cosocket.ae"
local anet = require "cosocket.anet"

local _M={}
local AE_READABLE=1
local AE_WRITABLE=2

local function socket_pool_gc(tab)
	local clear_map = {}
	for fd, s in pairs(tab) do
	  if s ~= nil then
		  if clear_map[s.aefd] == nil then
			  clear_map[s.aefd]={}
		  end
		  table.insert(clear_map[s.aefd],fd)
	  end
	end
	for k,v in pairs(clear_map) do
	  for i=1,#v do
		  ae.del(k,v[i])
		  anet.close(v[i])
	  end
	  ae.close(k)
	end
end

local socket_pool = setmetatable({},{__gc = socket_pool_gc}) -- for lua 5.1, __gc not work automatically, so you should call clear explictly

local main_co

local function int_endian_convert(str, start, size, flag)
    if str == nil or start+size > str:len() or size <=0 or start<1 then
        return nil
    end
    local res = 0
    for i=1,size,1 do
        local b = 0
        if flag>0 then
            b = string.byte(str,start+i-1)
        else
            b = string.byte(str,start+size-i)
        end
        if b == nil then
            return nil
        end
        res = res*256+b
    end
    return res
end
    

function _M.new_poll()
    if socket_pool == nil then
	    socket_pool = setmetatable({},{__gc = socket_pool_gc})
	end
    main_co = coroutine.running()
    return ae.create()
end
function _M.clear()
    socket_pool_gc(socket_pool)
    socket_pool=nil
return
end
local function IS_MAIN_CO()
    return assert(main_co)==coroutine.running()
end

_M.IS_MAIN_CO = IS_MAIN_CO

function _M.free_poll(aefd)
    for k,v in pairs(socket_pool) do
	    if v ~= nil and v.aefd == aefd then
		    ae.del(aefd,k)
		    anet.close(k)
		end
	end
    return ae.close(aefd)
end

local function ev_listen_handler(s,readable,writeable)
    -- print("ev_listen_handler",s.fd,readable,writeable)
    if readable and not writeable then
        coroutine.resume(s.co)
    end
end

function _M.listen(aefd,endpoint,on_accept)
    local host,port = endpoint:match("([^:]+):(.+)$")
    print("listen:",host,port)
    port = tonumber(port)
    local fd = anet.listen(host,port,32)
    local ret = anet.noblock(fd,1)
    local co = coroutine.create(function ()
               while true do
                   local clientfd,ip,clientport = anet.accept(fd)
                   if clientfd <= 0 or clientfd == nil then 
                       coroutine.yield(coroutine.running())
                   else
                       --anet.noblock(clientfd,1)
                       on_accept(clientfd,ip,clientport)
                   end
               end
           end)
    local s = {aefd=aefd,fd=fd,co=co,ev_handler = ev_listen_handler}
    socket_pool[fd] = s
    ae.add_read(aefd,fd)
    return fd
end

local function ev_client_handler(s,readable,writeable)
    assert(s)
    print("ev_client_handler: fd=", s.fd)
    local res = nil
    local is_resumed = false
    if readable then
        local n,buf = anet.read(s.fd,1024)
        if n<=0 then
            -- close
            print("client disconnect")
            ae.del(s.aefd,s.fd)
            anet.close(s.fd)
            socket_pool[s.fd]=nil
        elseif n>0 then
            -- s.read_buffer = s.read_buffer..buf
            s.read_buffer = table.concat({s.read_buffer,buf}, "")
            if s.prototype == 1 then
                -- s.read_need indicates end str
                local i,j=s.read_buffer:find(s.protoflag,1,true)
                if i~=nil then
                    local sta=coroutine.status(s.co)
                    if sta == "dead" then
                        is_resumed = false
                        ae.del(s.aefd,s.fd)
                        anet.close(s.fd)
                        socket_pool[s.fd]=nil
                        print(" coroutine for "..s.fd.." is dead before resume.")  
                    else
                        sta,res=coroutine.resume(s.co,i,j)
                        is_resumed = true
                    end
                    if is_resumed and sta == false then 
                          ae.del(s.aefd,s.fd)
                          anet.close(s.fd)
                          socket_pool[s.fd]=nil
                          print(" coroutine for "..s.fd.." is dead after resume")                     
		             end
                end
            elseif s.prototype == 2 then
                local msg_len = s.read_buffer:len()
                if msg_len >= s.protoflag then
                    local sta=coroutine.status(s.co)
                    if sta == "dead" then
                        is_resumed = false
                        ae.del(s.aefd,s.fd)
                        anet.close(s.fd)
                        socket_pool[s.fd]=nil
                        print(" coroutine for "..s.fd.." is dead before resume.")  
                    else
                        sta,res=coroutine.resume(s.co,s.protoflag)
                        is_resumed = true
                    end
                    if is_resumed and sta == false then 
                          ae.del(s.aefd,s.fd)
                          anet.close(s.fd)
                          socket_pool[s.fd]=nil
                          print(" coroutine for "..s.fd.." is dead after resume")                     
		             end
                  end
            elseif s.prototype == 3 then
                local msg_len = s.read_buffer:len()
                if msg_len >= s.protolensize then
                    local packet_length = 1024
                    if s.protoflag then
                        -- big endian
                        packet_length = int_endian_convert(s.read_buffer,1,s.protolensize,1)
                    else
                        -- little endian
                        packet_length = int_endian_convert(s.read_buffer,1,s.protolensize,0)
                    end
                    if packet_length == nil then
                        -- message error
                        s.read_buffer=""
                    end
                    print(s.read_buffer,msg_len,packet_length)
                    if msg_len >= packet_length then
                        local sta=coroutine.status(s.co)
                        if sta == "dead" then
                            is_resumed = false
                            ae.del(s.aefd,s.fd)
                            anet.close(s.fd)
                            socket_pool[s.fd]=nil
                            print(" coroutine for "..s.fd.." is dead before resume.")  
                        else
                            sta,res=coroutine.resume(s.co,packet_length)
                            is_resumed = true
                        end
                        if is_resumed and sta == false then 
                            ae.del(s.aefd,s.fd)
                            anet.close(s.fd)
                            socket_pool[s.fd]=nil
                            print(" coroutine for "..s.fd.." is dead after resume")                     
                        end
                    else
                        --print("required length=",packet_length)
                    end
                end
            else
                -- unknown prototype
                s.read_buffer=""
            end
        end
    end
    if writeable then
        if #s.write_buffer > 0 then
            local n = anet.write(s.fd,s.write_buffer)
            if n>0 then
                s.write_buffer = s.write_buffer:sub(n+1)
                if #s.write_buffer > 0 then
                    ae.enable(s.aefd,fd,readable,true)
                end
            else
                ae.del(s.aefd,s.fd)
                anet.close(s.fd)
                socket_pool[s.fd]=nil
            end
        else
            ae.enable(s.aefd,fd,readable,false)
        end
    end
    if not readable and not writeable then
        ae.del(s.aefd,s.fd)
        anet.close(s.fd)
        socket_pool[s.fd]=nil
    end
end

-- proto_type: 1 for string, 2 for fixed length, 3 for dynamic bytesstream
-- if proto_type == 1, param flag must be set a valid string (spe:len()>0)
-- fi proto_type == 2, param flag must be set a valid number (size>0)
-- fi proto_type == 3, param flag must be set a boolean, true for bigendian otherwise littlendian

local function start(aefd,fd,func,prototype,flag,size)
    --[[if aefd==nil or fd == nil or func == nil then
        return -1
    end]]
    
    
    local co = coroutine.create(function()
                   func(fd)
               end)
    local s = socket_pool[fd]
    local ret
    if s~=nil then
        ret = ae.enable(aefd,fd,true,false)
    else
        ret = ae.add_read(aefd,fd)
    end
    assert(ret==0)
    assert(type(prototype)=="number")
    assert(prototype>0 and prototype<4)
    if prototype == 1 then
        assert(type(flag)=="string" and flag:len()>0)
    elseif prototype == 2 then
       assert(type(flag)=="number" and flag > 0)
    elseif prototype == 3 then
       assert(type(flag)=="boolean")
       assert(type(size)=="number" and (size == 2 or size == 4 or size == 8))
    else
       assert(false)
    end
    socket_pool[fd] = nil   -- delete original dict
    socket_pool[fd] = {aefd=aefd,fd=fd,co=co,prototype=prototype,protoflag=flag,protolensize=size,
                       read_buffer="",write_buffer="",
                       ev_handler=ev_client_handler}
    local sta, res = coroutine.resume(co)
    if sta == false then
        ae.del(aefd,fd)
        anet.close(fd)
        socket_pool[fd] = nil
    end
    return 0
end

_M.start = start

function _M.readpacket(fd)
    local s = assert(socket_pool[fd])
    if s.prototype == 1 then
        local i,j = s.read_buffer:find(s.protoflag,1,true)
        if i == nil then
            i,j = coroutine.yield(s.co)
        end
        local buf = s.read_buffer:sub(1,i-1)
        s.read_buffer = s.read_buffer:sub(j+1)
        return buf
    elseif s.prototype == 2 then
        local msg_len = s.read_buffer:len()
        if msg_len < s.protoflag then
            msg_len = coroutine.yield(s.co)
        end
        local buf = s.read_buffer:sub(1,msg_len)
        s.read_buffer = s.read_buffer:sub(msg_len+1,-1)
        return buf
    elseif s.prototype == 3 then
        assert(s.protolensize==2 or s.protolensize == 4 or s.protolensize == 8)
        local msg_len = s.read_buffer:len()
        if msg_len < 1 then
            msg_len = coroutine.yield(s.co)
        end
        local endianess = -1
        if s.protoflag then
            endianess = 1
        end
        local sz = int_endian_convert(s.read_buffer,1,s.protolensize,endianess)
        if sz ~= nil and sz <= msg_len then
            local buf = s.read_buffer:sub(1,sz)
            s.read_buffer = s.read_buffer:sub(sz+1,-1)
            return buf
        end
        msg_len = coroutine.yield(s.co)
        local buf = s.read_buffer:sub(1,msg_len)
        s.read_buffer = s.read_buffer:sub(msg_len+1,-1)
        return buf
    end
end

function _M.write(fd,buf)
    local s = assert(socket_pool[fd])
    s.write_buffer = s.write_buffer..buf
    local ret,err = anet.write(fd,s.write_buffer)
    if ret<0 then
        -- disconnect
        print("write:",err)
		ae.del(s.aefd,s.fd)
        anet.close(s.fd)
        socket_pool[s.fd]=nil
    elseif ret>0 then
        s.write_buffer = s.write_buffer:sub(ret+1)
		if #s.write_buffer > 0 then
		    ae.enable(s.aefd,fd,true,false)
		end
    end
    if ret == 0 and #s.write_buffer > 0 then
        -- enable write event
        ae.enable(s.aefd,fd,true,true)
    end
end

local function ev_connect_handler(s,readable,writeable)
     if writeable and not readable then
         coroutine.resume(s.co)
     else
         ae.del(s.aefd,s.fd)
         anet.close(s.fd)
         socket_pool[s.fd]=nil
     end
end

function _M.block_connect(ip,port)
    local fd = anet.connect(ip,port,-1)
    if fd < 0 then
        -- connect failed
        return -1
    end
    return fd
end

function _M.connect(aefd,endpoint,func)
    local host,port = endpoint:match("([^:]+):(.+)$")
    local fd = anet.connect(host,port,0)
    if fd > 0 then
        local co=coroutine.create(function ()
            local s = socket_pool[fd]
            assert(s)
            ae.enable(aefd,fd,true,false)
            func(fd)
        end)
        ae.add_write(aefd,fd)
        socket_pool[fd] = {aefd=aefd,fd=fd,ev_handler=ev_connect_handler,co=co}
        return fd
    end
    return -1
end

local function event_handler(fd,readable,writeable)
    local s = assert(socket_pool[fd])
    s:ev_handler(readable,writeable)
end

function _M.run(aefd,timeout)
    while true do
        local ret = ae.poll(aefd,timeout or -1,event_handler)
        if ret<0 then
            break;
        end
    end
end

return _M
